package com.liuqi.openai.core.client.httpclient;

import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.PropertyNamingStrategy;
import com.alibaba.fastjson.serializer.SerializeConfig;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.liuqi.openai.core.client.Executor;
import com.liuqi.openai.core.client.OpenAiRemoteException;
import com.liuqi.openai.core.client.StreamHandler;
import com.liuqi.openai.core.merge.MergeChunk;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.function.Supplier;
import lombok.extern.slf4j.Slf4j;
import org.apache.hc.client5.http.classic.HttpClient;
import org.apache.hc.client5.http.classic.methods.HttpPost;
import org.apache.hc.client5.http.classic.methods.HttpUriRequest;
import org.apache.hc.core5.http.ContentType;
import org.apache.hc.core5.http.HttpEntity;
import org.apache.hc.core5.http.HttpStatus;
import org.apache.hc.core5.http.io.entity.EntityUtils;
import org.apache.hc.core5.http.io.entity.StringEntity;

/**
 * @author liuqi
 * @date 2025/7/18
 **/
@Slf4j
public class HttpClientExecutor<Request, Response> implements Executor<Response> {

    /**
     * fast json 序列化, 驼峰转下划线
     */
    final SerializeConfig serializeConfig = new SerializeConfig();
    {
        serializeConfig.setPropertyNamingStrategy(PropertyNamingStrategy.SnakeCase);
    }

    /**
     * httpclient 实例
     */
    private final HttpClient httpClient;

    /**
     * 请求地址
     */
    private final String endpointUrl;

    /**
     * 请求数据
     */
    private final Supplier<Request> request;

    /**
     * 响应解码类
     */
    private final Class<Response> responseClass;

    /**
     * 是否流式响应
     */
    private final boolean streaming;

    /**
     * 流式响应回调处理器
     *
     * 如果此实例不为空, 则应用流式读取数据
     */
    private StreamHandler<Response> streamHandler;

    /**
     * http post 请求, 全局保存, 适用与关闭连接中断响应
     */
    private HttpUriRequest httpUriRequest;

    public HttpClientExecutor(HttpClient httpClient,
                              String endpointUrl,
                              Supplier<Request> request,
                              Class<Response> responseClass) {
        this(httpClient, endpointUrl, false, request, responseClass);
    }

    public HttpClientExecutor(HttpClient httpClient,
                              String endpointUrl,
                              boolean streaming,
                              Supplier<Request> request,
                              Class<Response> responseClass) {
        this.httpClient = httpClient;
        this.endpointUrl = endpointUrl;
        this.streaming = streaming;
        this.request = request;
        this.responseClass = responseClass;
    }

    @Override
    public Response execute() throws OpenAiRemoteException {
        try {
            httpUriRequest = new HttpPost(endpointUrl);

            String body = JSONObject.toJSONString(request.get(), serializeConfig, SerializerFeature.WriteEnumUsingToString);
            if (log.isTraceEnabled()) {
                log.trace("OpanAi HttpClientExecutor, request body: " + body);
            }

            StringEntity requestEntity = new StringEntity(
                    body,
                    ContentType.APPLICATION_JSON);
            httpUriRequest.setEntity(requestEntity);

            return httpClient.execute(httpUriRequest, response -> {
                // 状态码
                int code = response.getCode();
                // 状态码不等于 200 一律抛出异常
                if (code != HttpStatus.SC_OK) {
                    throw new IOException("responseCode = " + code + ", errorResponse = " + EntityUtils.toString(response.getEntity()));
                }

                // 非流式
                if (streaming) {
                    // 流式
                    HttpEntity entity = response.getEntity();
                    if (entity != null) {
                        return processSSEStream(entity.getContent());
                    }
                    throw new IOException("responseCode = " + code + ", errorResponse = null");
                } else {
                    return JSONObject.parseObject(EntityUtils.toString(response.getEntity()), responseClass);
                }
            });
        } catch (Throwable t) {
            if (streamHandler != null) {
                streamHandler.onError(t);
            }

            if (t instanceof  OpenAiRemoteException) {
                throw (OpenAiRemoteException) t;
            } else {
                throw new OpenAiRemoteException("OpenAiApi execute error. ", t);
            }
        }
    }

    @Override
    public Executor<Response> injectStreamHandler(StreamHandler<Response> streamHandler) {
        if (this.streamHandler != null) {
            throw new RuntimeException("HttpClientExecutor current property streamHandler is not empty. ");
        }
        this.streamHandler = streamHandler;
        return this;
    }

    @Override
    public void close() {
        if (httpUriRequest != null && !httpUriRequest.isAborted()) {
            httpUriRequest.abort();
        }
    }

    /**
     * 读取流式响应
     *
     * @param inputStream
     * @return
     */
    private Response processSSEStream(InputStream inputStream) throws IOException {
        MergeChunk mergeChunk = MergeChunk.getMergeChunk(responseClass);

        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
            String line;
            while ((line = reader.readLine()) != null) {
                if (line.isEmpty() || line.startsWith(":")) {
                    continue;
                }
                int colonIndex = line.indexOf(":");
                if (colonIndex < 0) {
                    continue;
                }
                String field = line.substring(0, colonIndex).trim();
                String value = colonIndex < line.length() - 1 ?
                        line.substring(colonIndex + 1).trim() : "";

                // 只解析 data: 的数据
                if ("data".equals(field)) {
                    if ("[DONE]".equals(value)) {
                        // done
                        if (streamHandler != null) {
                            streamHandler.onComplete();
                        }
                    } else {
                        // json 序列化
                        Response response = JSONObject.parseObject(value, responseClass);

                        // 通知流式回调接口, 接收到了响应
                        if (streamHandler != null) {
                            streamHandler.onResponse(response);
                        }

                        // 合并数据块
                        mergeChunk.merge(response);
                    }
                }
            }

            // 返回合并后的数据块
            return (Response) mergeChunk.getResult();
        } catch (IOException e) {
            throw e;
        }
    }

}
